import os
os.environ['WANDB_DISABLED'] = 'true'
os.environ['OMP_NUM_THREADS'] = '2'
os.environ['CUDA_VISIBLE_DEVICES'] = '2'  # 指定可见的GPU
from ultralytics import YOLO
from ultralytics.data import build_dataloader, build_llltData
from ultralytics.cfg import get_cfg
from ultralytics.models.yolo.detect import DetectionTrainer, llltTrainer
import argparse
import yaml
import torch.distributed as dist
import torch
from ultralytics.utils.torch_utils import select_device
from ultralytics.models.yolo.detect import DetectionValidator, lltValidator



# 参数解析
parser = argparse.ArgumentParser()
parser.add_argument("--data", type=str, default="lllt.yaml", help="data yaml file")
parser.add_argument("--model", type=str, default="snn_yolov8l.yaml", help="model yaml file")
parser.add_argument("--weights", type=str, default="./runs/train/exp_t_yolov8m_1/weights/best.pt", help="weights file")
parser.add_argument("--batch", type=int, default=4, help="batch size")
parser.add_argument("--workers", type=int, default=16, help="workers num")
args = parser.parse_args()

# 主函数
def main():
    # 初始化分布式训练
    # is_distributed, rank, world_size, gpu = init_distributed_mode()
    # print(gpu)
    # 加载数据配置
    # with open("./ultralytics/cfg/datasets/"+ args.data, "r") as file:
    #     data = yaml.safe_load(file)
    # image_path = data.get('path')

    # 训练参数
    val_args = dict(
        model=args.weights,
        data=args.data,
        batch=args.batch,
        device=[2], 
        # half=True,
        # amp=False,
        # single_cls=False,
        # optimizer='auto',
        project='runs/detect',
        name='test_rescue_yolov8m',
        exist_ok=True,
        # resume=True,
    )

    # args = dict(model='yolov8n.pt', data='coco8.yaml')
    validator = lltValidator(args=val_args, json_path="/data1/lkf24/data/NESR/annotations")
    validator()

if __name__ == "__main__":
    main()